import os
import cv2
import numpy as np
import json
import copy
from tqdm import tqdm
import base64
import random

def get_img_list(path):
    img_list = []
    for root, dirs, files in os.walk(path):
        for f in files:
            if f.endswith('.jpg'):
                img_list.append(os.path.join(root, f))

    assert len(img_list) > 0, "没有图片！"
    return img_list


def main():
    file_dir = r"data/ori_data/4.19"
    save_dir = r"data/ori_data/4.19_crop"

    img_list = get_img_list(file_dir)

    for img_p in tqdm(img_list):
        json_p = img_p.replace('.jpg', '.json')
        img = cv2.imread(img_p)
        img_h, img_w, _ = img.shape
        with open(json_p, 'r') as f:
            content = json.loads(f.read())

        contours = []
        for shape in content['shapes']:
            points = shape['points']
            contours.append(np.array(points, dtype=np.float32))
        
        total_points = np.concatenate(contours, 0)
        x0 = np.min(total_points[:, 0])
        x1 = np.max(total_points[:, 0])
        y0 = np.min(total_points[:, 1])
        y1 = np.max(total_points[:, 1])
        w = x1 - x0
        h = y1 - y0

        if w >= 500:
            new_w = w
        else:
            new_w = 500
            # x0 = max(0, x0 - 250)
            x0 = max(0, x0 - random.randint(0, 500 - int(w)))

        if h >= 500:
            new_h = h
        else:
            new_h = 500
            # y0 = max(0, y0 - 250)
            y0 = max(0, y0 - random.randint(0, 500 - int(h)))

        new_x0 = int(max(0, x0 - 50))
        new_x1 = int(min(img_w, new_x0 + new_w + 100))
        new_y0 = int(max(0, y0 - 50))
        new_y1 = int(min(img_h, new_y0 + new_h + 100))

        new_img = img[new_y0:new_y1, new_x0:new_x1, :]

        new_shapes = []
        shapes = content['shapes']
        for shape in shapes:
            points = shape['points']
            points_arr = np.array(points, dtype=np.float32)
            points_arr[:, 0] -= new_x0
            points_arr[:, 1] -= new_y0
            new_shape = shape.copy()
            new_shape['points'] = points_arr.tolist()
            new_shapes.append(new_shape)
        
        new_content = content.copy()
        new_content['shapes'] = new_shapes
        new_content['imageHeight'] = new_img.shape[0]
        new_content['imageWidth'] = new_img.shape[1]
        img_str = cv2.imencode('.jpg', new_img)[1].tobytes()  # 将图片编码成流数据，放到内存缓存中，然后转化成string格式
        b64_code = base64.b64encode(img_str).decode() # 编码成base64
        new_content['imageData'] = b64_code
        
        save_img_path = os.path.join(save_dir, new_content['imagePath'])
        save_json_path = save_img_path.replace('.jpg', '.json')

        cv2.imwrite(save_img_path, new_img)
        with open(save_json_path, 'w') as f:
            json.dump(new_content, f, indent=4)


        _ = 1




if __name__ == "__main__":
    main()
